{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fast LLM reasoning with DeepSeek-R1-Distill-Llama-8B and FastDraft\n",
    "\n",
    "[DeepSeek-R1](https://github.com/deepseek-ai/DeepSeek-R1/blob/main/DeepSeek_R1.pdf) is an open-source reasoning model developed by DeepSeek to address tasks requiring logical inference, mathematical problem-solving, and real-time decision-making. With DeepSeek-R1, you can follow its logic, making it easier to understand and, if necessary, challenge its output. \n",
    "This capability gives reasoning models an edge in fields where outcomes need to be explainable, like research or complex decision-making.\n",
    "Distillation in AI creates smaller, more efficient models from larger ones, preserving much of their reasoning power while reducing computational demands. DeepSeek applied this technique to create a suite of distilled models from R1, using Qwen and Llama architectures. \n",
    "That allows us to try DeepSeek-R1 capability locally on usual laptops.<br>\n",
    "Please check [this](https://github.com/openvinotoolkit/openvino_notebooks/tree/latest/notebooks/deepseek-r1) great tutorial for running DeepSeek-R1 distilled models on your laptop with OpenVINO.\n",
    "\n",
    "In this notebook we will show you how to speedup DeepSeek-R1 inference on your laptop with FastDraft, our latest publication.\n",
    "\n",
    "## Intel Labs' FastDraft\n",
    "FastDraft is a novel and efficient approach for pre-training and aligning a draft model to any LLM to be used with speculative decoding, by incorporating efficient pre-training followed by fine-tuning over synthetic datasets generated by the target model. \n",
    "FastDraft was presented in a [paper](https://arxiv.org/abs/2411.11055) at [ENLSP@NeurIPS24](https://neurips2024-enlsp.github.io/accepted_papers.html) by Intel Labs.\n",
    "FastDraft pre-trained draft models achieve impressive results in key metrics of acceptance rate, block efficiency and up to 3x memory bound speed up when evaluated on code completion and up to 2x in summarization, text completion and instruction tasks and unlock large language models inference on AI-PC and other edge-devices.\n",
    "\n",
    "In this notebook we will use the Llama-3.1 FastDraft model created as part of the project to accelerate [DeepSeek-R1-Distill-Llama-8B](https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Llama-8B). Since this model shares the same vocabulary as the Llama-3.1 family, our Llama-3.1 FastDraft is compatible with DeepSeek-R1-Distill-Llama-8B model.\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "\n",
    "> Note: we recommend running this notebook in a virtual environment. \n",
    "\n",
    "Install required dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import platform\n",
    "\n",
    "os.environ[\"GIT_CLONE_PROTECTION_ACTIVE\"] = \"false\"\n",
    "\n",
    "%pip install -q -U --pre \"openvino==2025.1.0.dev20250120\" \"openvino-tokenizers==2025.1.0.dev20250120\" \"openvino-genai==2025.1.0.dev20250120\" --extra-index-url https://storage.openvinotoolkit.org/simple/wheels/nightly\n",
    "%pip install -q --extra-index-url https://download.pytorch.org/whl/cpu \\\n",
    "\"optimum-intel==1.20.1\" \\\n",
    "\"optimum==1.23.3\" \\\n",
    "\"nncf==2.14.1\" \\\n",
    "\"torch>=2.1\" \\\n",
    "\"datasets<4.0.0\" \\\n",
    "\"accelerate\" \\\n",
    "\"gradio>=4.19,<6\" \\\n",
    "\"transformers==4.46.3\" \\\n",
    "\"einops\" \"tiktoken\" \\\n",
    "\"numpy==2.0.2\"\n",
    "\n",
    "if platform.system() == \"Darwin\":\n",
    "    %pip install \"numpy<2.0.0\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generation with OpenVINO GenAI\n",
    "We will first need to convert the model to OpenVINO format and then we will be ready to run the model on your laptop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "model_name = \"DeepSeek-R1-Distill-Llama-8B\"\n",
    "model_id = \"deepseek-ai/DeepSeek-R1-Distill-Llama-8B\"\n",
    "model_dir = Path(\"./DeepSeek-R1-Distill-Llama-8B-int4-ov\")\n",
    "device = \"GPU\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Convert model using Optimum-CLI tool\n",
    "\n",
    "\n",
    "\n",
    "🤗 [Optimum Intel](https://huggingface.co/docs/optimum/intel/index) is the interface between the 🤗 [Transformers](https://huggingface.co/docs/transformers/index) and OpenVINO to accelerate end-to-end pipelines on Intel architectures. It provides ease-to-use cli interface for exporting models to [OpenVINO Intermediate Representation (IR)](https://docs.openvino.ai/2024/documentation/openvino-ir-format.html) format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not model_dir.exists():\n",
    "    ! optimum-cli export openvino --model $model_id --task text-generation-with-past --weight-format int4 $model_dir\n",
    "\n",
    "# convert OV tokenizer if needed\n",
    "if not (model_dir / \"openvino_tokenizer.xml\").exists():\n",
    "    ! convert_tokenizer $model_dir --with-detokenizer -o $model_dir"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate pipeline with OpenVINO Generate API\n",
    "\n",
    "We will use [OpenVINO Generate API](https://github.com/openvinotoolkit/openvino.genai/blob/master/src/README.md) to create pipelines to run an inference with OpenVINO Runtime. \n",
    "\n",
    "Firstly we need to create a pipeline with `LLMPipeline`. `LLMPipeline` is the main object used for text generation using LLM in OpenVINO GenAI API. You can construct it straight away from the folder with the converted model. We will provide directory with model and device for `LLMPipeline`. Then we run `generate` method and get the output in text format.\n",
    "Additionally, we can configure parameters for decoding. We can create the default config with `ov_genai.GenerationConfig()`, setup parameters, and apply the updated version with `set_generation_config(config)` or put config directly to `generate()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openvino_genai as ov_genai\n",
    "\n",
    "from llm_pipeline_with_hf_tokenizer import LLMPipelineWithHFTokenizer\n",
    "\n",
    "\n",
    "def streamer(subword):\n",
    "    print(subword, end=\"\", flush=True)\n",
    "    return False\n",
    "\n",
    "\n",
    "# Define scheduler\n",
    "scheduler_config = ov_genai.SchedulerConfig()\n",
    "scheduler_config.num_kv_blocks = 2048 // 16\n",
    "scheduler_config.dynamic_split_fuse = False\n",
    "scheduler_config.max_num_batched_tokens = 2048\n",
    "\n",
    "pipe = LLMPipelineWithHFTokenizer(model_dir, device, scheduler_config=scheduler_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After instantiating the pipeline we are ready to generate with the model.\n",
    "\n",
    "DeepSeek models work best when called with their custom chat template so we pass the query to the model in a chat format with `apply_chat_template=True` argument.\n",
    "\n",
    "To get a more accurate measurement of the generation time, we do a warmup step to let the model allocate memory and compile any kernels it needs to reach its full potential."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "\n",
    "generation_config = ov_genai.GenerationConfig()\n",
    "\n",
    "input_prompt = [{\"role\": \"user\", \"content\": \"Which number is bigger 9.11 or 9.9?\"}]\n",
    "\n",
    "# We will first do a short warmup to the model so the time measurement will not include the warmup overhead.\n",
    "generation_config.max_new_tokens = 8\n",
    "pipe.generate(input_prompt, generation_config, apply_chat_template=True)\n",
    "\n",
    "# Now we can measure the time and see the result\n",
    "generation_config.max_new_tokens = 1024\n",
    "\n",
    "start = time.perf_counter()\n",
    "result = pipe.generate(input_prompt, generation_config, streamer, apply_chat_template=True)\n",
    "ar_gen_time = time.perf_counter() - start"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gc\n",
    "\n",
    "print(f\"Generation took {ar_gen_time:.3f} seconds\")\n",
    "del pipe\n",
    "gc.collect()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Acceleration with FastDraft and speculative decoding\n",
    "Speculative decoding is a lossless decoding paradigm introduced in a recent [ICML paper](https://arxiv.org/abs/2211.17192) for accelerating auto-regressive generation with LLMs.\n",
    "The method aims to mitigate the inherent latency bottleneck caused by the sequential nature of auto-regressive generation.\n",
    "Speculative decoding employs a draft language model to generate a block of \\(\\gamma\\) candidate tokens.\n",
    "The LLM, referred to as the target model, then processes these candidate tokens in parallel.\n",
    "The algorithm examines each token's probability distribution, calculated by both the target and draft models, to determine whether the token should be accepted or rejected.\n",
    "\n",
    "In this section we will show how to accelerate the generation of the DeepSeek using our Llama-3.1 FastDraft and OpenVINO GenAI speculative decoding pipeline.\n",
    "\n",
    "First we will download our draft and then we will initialize a generation pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import filecmp\n",
    "import shutil\n",
    "\n",
    "import huggingface_hub as hf_hub\n",
    "\n",
    "draft_model_id = \"OpenVINO/Llama-3.1-8B-Instruct-FastDraft-150M-int8-ov\"\n",
    "draft_model_path = Path(\"Llama_3.1_FastDraft\")\n",
    "\n",
    "if not draft_model_path.exists():\n",
    "    hf_hub.snapshot_download(draft_model_id, local_dir=draft_model_path)\n",
    "\n",
    "# We need tokenizers to match between the target and draft model so we apply this workaround\n",
    "if not filecmp.cmp(str(model_dir / \"openvino_tokenizer.xml\"), str(draft_model_path / \"openvino_tokenizer.xml\"), shallow=False):\n",
    "    for fname in [\"openvino_tokenizer.xml\", \"openvino_tokenizer.bin\", \"openvino_detokenizer.xml\", \"openvino_detokenizer.bin\"]:\n",
    "        shutil.copy(model_dir / fname, draft_model_path / fname)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define schedulers\n",
    "scheduler_config = ov_genai.SchedulerConfig()\n",
    "scheduler_config.num_kv_blocks = 2048 // 16\n",
    "scheduler_config.dynamic_split_fuse = False\n",
    "scheduler_config.max_num_batched_tokens = 2048\n",
    "\n",
    "draft_scheduler_config = ov_genai.SchedulerConfig()\n",
    "draft_scheduler_config.num_kv_blocks = 2048 // 16\n",
    "draft_scheduler_config.dynamic_split_fuse = False\n",
    "draft_scheduler_config.max_num_batched_tokens = 2048\n",
    "\n",
    "draft_model = ov_genai.draft_model(draft_model_path, device, scheduler_config=draft_scheduler_config)\n",
    "\n",
    "pipe = LLMPipelineWithHFTokenizer(model_dir, device, scheduler_config=scheduler_config, draft_model=draft_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we are ready to generate with our speculative decoding pipeline, as in the previous section, we will do a small warmup step before measuring the generation time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We need to define in the generation config how many tokens the draft should predict in each cycle\n",
    "generation_config.num_assistant_tokens = 3\n",
    "\n",
    "# Again we will do a short warmup before measuring time for the model\n",
    "generation_config.max_new_tokens = 8\n",
    "pipe.generate(input_prompt, generation_config, apply_chat_template=True)\n",
    "\n",
    "# Now we can measure the time and see the result\n",
    "generation_config.max_new_tokens = 1024\n",
    "\n",
    "start = time.perf_counter()\n",
    "result = pipe.generate(input_prompt, generation_config, streamer, apply_chat_template=True)\n",
    "sd_gen_time = time.perf_counter() - start"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Generation took {sd_gen_time:.3f} seconds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In our FastDraft paper we have measured a 1.5x speedup in average with our draft after aligning it to the target model, in this case the draft wasn't align to the target mode, however, we still saw roughly the same speedup results.\n",
    "Let's calculate the speedup results for the specific example we used:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"End to end speedup with FastDraft and speculative decoding is {ar_gen_time / sd_gen_time:.2f}x\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate Speculative Decoding Speedup On Multiple Examples\n",
    "\n",
    "In this section we compare auto-regressive generation and speculative-decoding generation with DeepSeek-R1-Distill-Llama-8B model on multiple examples. \n",
    "We use ~50 example-prompts taken from [MT-Bench](https://huggingface.co/datasets/HuggingFaceH4/mt_bench_prompts) and from [databricks-dolly-15k](https://huggingface.co/datasets/databricks/databricks-dolly-15k) datasets.\n",
    "We loop over these examples and measure generation times, first without speculative-decoding and later with speculative-decoding. Eventually we compare generation times for both methods and compute the average speedup gain."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Run target model without speculative decoding\n",
    "As in previous section, we will first run generation without speculative-decoding, but this time we will run it over ~50 examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<frozen importlib.util>:207: DeprecationWarning: The `openvino.runtime` module is deprecated and will be removed in the 2026.0 release. Please replace `openvino.runtime` with `openvino`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading model from DeepSeek-R1-Distill-Llama-8B-int4-ov\n",
      "Loading prompts...\n",
      "Running Auto-Regressive generation...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████| 53/53 [21:14<00:00, 24.05s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Done\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "211"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import openvino_genai as ov_genai\n",
    "import sys\n",
    "import time\n",
    "from tqdm import tqdm\n",
    "from llm_pipeline_with_hf_tokenizer import LLMPipelineWithHFTokenizer\n",
    "\n",
    "print(f\"Loading model from {model_dir}\")\n",
    "\n",
    "# Define scheduler\n",
    "scheduler_config = ov_genai.SchedulerConfig()\n",
    "scheduler_config.num_kv_blocks = 2048 // 16\n",
    "scheduler_config.dynamic_split_fuse = False\n",
    "scheduler_config.max_num_batched_tokens = 2048\n",
    "\n",
    "pipe = LLMPipelineWithHFTokenizer(model_dir, device, scheduler_config=scheduler_config)\n",
    "\n",
    "generation_config = ov_genai.GenerationConfig()\n",
    "\n",
    "\n",
    "print(\"Loading prompts...\")\n",
    "import json\n",
    "\n",
    "f = open(\"prompts.json\")\n",
    "prompts = json.load(f)\n",
    "prompts = [[{\"role\": \"user\", \"content\": p}] for p in prompts]\n",
    "\n",
    "# We will first do a short warmup to the model so the time measurement will not include the warmup overhead.\n",
    "generation_config.max_new_tokens = 8\n",
    "pipe.generate(prompts[0], generation_config, apply_chat_template=True)\n",
    "\n",
    "# finished warmup step, let's run our examples\n",
    "generation_config.max_new_tokens = 2048\n",
    "times_auto_regressive = []\n",
    "print(\"Running Auto-Regressive generation...\")\n",
    "for prompt in tqdm(prompts):\n",
    "    start_time = time.perf_counter()\n",
    "    result = pipe.generate(prompt, generation_config, apply_chat_template=True)\n",
    "    end_time = time.perf_counter()\n",
    "    times_auto_regressive.append(end_time - start_time)\n",
    "print(\"Done\")\n",
    "\n",
    "import gc\n",
    "\n",
    "del pipe\n",
    "gc.collect()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Run target model with speculative decoding\n",
    "Now we will run generation with speculative-decoding over the same ~50 examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading prompts...\n",
      "Running Speculative Decoding generation...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  4%|███▏                                                                               | 2/53 [00:15<06:27,  7.60s/it]"
     ]
    }
   ],
   "source": [
    "import openvino_genai as ov_genai\n",
    "import time\n",
    "from tqdm import tqdm\n",
    "from llm_pipeline_with_hf_tokenizer import LLMPipelineWithHFTokenizer\n",
    "\n",
    "\n",
    "print(\"Loading prompts...\")\n",
    "import json\n",
    "\n",
    "f = open(\"prompts.json\")\n",
    "prompts = json.load(f)\n",
    "prompts = [[{\"role\": \"user\", \"content\": p}] for p in prompts]\n",
    "\n",
    "# Define scheduler\n",
    "scheduler_config = ov_genai.SchedulerConfig()\n",
    "scheduler_config.num_kv_blocks = 2048 // 16\n",
    "scheduler_config.dynamic_split_fuse = False\n",
    "scheduler_config.max_num_batched_tokens = 2048\n",
    "\n",
    "# Define scheduler for the draft\n",
    "\n",
    "draft_scheduler_config = ov_genai.SchedulerConfig()\n",
    "draft_scheduler_config.num_kv_blocks = 2048 // 16\n",
    "draft_scheduler_config.dynamic_split_fuse = False\n",
    "draft_scheduler_config.max_num_batched_tokens = 2048\n",
    "scheduler_config.enable_prefix_caching = False\n",
    "scheduler_config.use_cache_eviction = False\n",
    "\n",
    "draft_model = ov_genai.draft_model(draft_model_path, device, scheduler_config=draft_scheduler_config)\n",
    "\n",
    "pipe = LLMPipelineWithHFTokenizer(model_dir, device, draft_model=draft_model, scheduler_config=scheduler_config)\n",
    "\n",
    "generation_config = ov_genai.GenerationConfig()\n",
    "generation_config.num_assistant_tokens = 3\n",
    "\n",
    "# Again, We will first do a short warmup\n",
    "generation_config.max_new_tokens = 8\n",
    "pipe.generate(prompts[0], generation_config, apply_chat_template=True)\n",
    "\n",
    "# finished warmup step, let's run our examples\n",
    "generation_config.max_new_tokens = 2048\n",
    "times_speculative_decoding = []\n",
    "\n",
    "print(\"Running Speculative Decoding generation...\")\n",
    "for prompt in tqdm(prompts):\n",
    "    start_time = time.perf_counter()\n",
    "    result = pipe.generate(prompt, generation_config, apply_chat_template=True)\n",
    "    end_time = time.perf_counter()\n",
    "    times_speculative_decoding.append(end_time - start_time)\n",
    "print(\"Done\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Calculate total speedup\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "avg_speedup = sum([x / y for x, y in zip(times_auto_regressive, times_speculative_decoding)]) / len(prompts)\n",
    "print(f\"average speedup: {avg_speedup:.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that by using speculative-decoding with FastDraft we can accelerate DeepSeek-R1-Distill-Llama-8B generation by ~1.4x on avarage.\n",
    "Please note that better performance is expected when using a python script to run this process."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run Chatbot\n",
    "\n",
    "Now, when model created, we can setup Chatbot interface using [Gradio](https://www.gradio.app/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from gradio_helper import make_demo\n",
    "\n",
    "stop_strings = [\"<｜end▁of▁sentence｜>\", \"<｜User｜>\", \"</User|>\", \"<|User|>\", \"<|end_of_sentence|>\", \"</｜\"]\n",
    "demo = make_demo(pipe, stop_strings, model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "demo.launch(inline=False, inbrowser=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# demo.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.8"
  },
  "openvino_notebooks": {
   "imageUrl": "https://user-images.githubusercontent.com/29454499/255799218-611e7189-8979-4ef5-8a80-5a75e0136b50.png",
   "tags": {
    "categories": [
     "Model Demos",
     "AI Trends"
    ],
    "libraries": [],
    "other": [
     "LLM"
    ],
    "tasks": [
     "Text Generation",
     "Conversational"
    ]
   }
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
